import logging
import os

from langchain_community.llms import LlamaCpp
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableSequence

os.environ["http_proxy"] = "http://127.0.0.1:7890"
os.environ["https_proxy"] = "http://127.0.0.1:7890"

# 配置日志系统，用于输出程序运行信息
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

MODEL_PATH = "F:/MODEL/lmstudio-community/Qwen3-8B-GGUF/Qwen3-8B-Q8_0.gguf"  # 本地模型路径
MAX_TOKENS = 512  # 增加生成长度以提供更完整的回答
N_THREADS = os.cpu_count()  # 利用CPU核心进行并行计算
CONTEXT_SIZE = 32768  # 上下文长度，模型能处理的最大上下文token数


def load_llm():
    logger.info("⏳ 正在加载模型...")
    try:
        llm = LlamaCpp(
            model_path=MODEL_PATH,  # 模型文件路径
            n_ctx=CONTEXT_SIZE,  # 上下文长度
            max_tokens=MAX_TOKENS,  # 最大生成token数
            temperature=0.3,  # 进一步降低随机性以提高稳定性
            n_gpu_layers=-1,  # 根据CUDA可用性自动设置GPU层
            n_threads=N_THREADS,  # CPU线程数
            n_batch=512,  # 批处理大小，提高处理效率
            verbose=False,  # 是否输出详细日志
            use_mlock=True,  # 锁定模型在RAM中以提高性能
            seed=42,  # 随机种子，保证结果可复现
            f16_kv=True,  # 使用半精度键值缓存，节省内存
            use_mmap=False,  # 禁用内存映射，避免相关错误
            repeat_penalty=1.1,  # 减少重复惩罚，改善输出质量 (参考项目经验)
            top_p=0.8,  # nucleus sampling参数
            top_k=30,  # 限制词汇选择范围
        )
        logger.info("✅ 模型加载成功")
        return llm
    except Exception as e:
        logger.error(f"❌ 模型加载失败: {str(e)}")
        raise  # 重新抛出异常以确保程序正确处理


def create_prompt_template():
    template = """
    你是一个专业的AI助手，专门帮助用户解答各种问题。

    回答要求：
    1. 使用简体中文回答
    2. 保持回答简洁明了，重点突出
    3. 如果不知道答案，请直接回答"我不知道"，不要编造内容
    4. 只输出一次答案，不要重复
    5. 回答长度控制在5句话以内

    用户问题：{question}

    请按照以下步骤思考并回答：
    1. 理解用户问题的核心要点
    2. 组织简洁准确的语言回答
    3. 检查回答是否符合要求

    你的回答："""

    return PromptTemplate(template=template, input_variables=["question"])


def create_qa_system(llm):
    print("⏳ 创建对话系统...")
    prompt = create_prompt_template()
    # 使用 RunnableSequence 替代 LLMChain
    chain = RunnableSequence(prompt, llm)
    return chain


def cuda_chat_main():
    try:
        qwen_llm = load_llm()
        chat_chain = create_qa_system(qwen_llm)
        print("创建对话系统成功。。")

        while True:
            try:
                # 获取用户输入
                user_input = input("\n你: ").strip()

                # 处理特殊命令 - 退出程序
                if user_input.lower() in ["/exit", "退出", "exit"]:
                    print("\n助手: 再见！感谢你的聊天~")
                    break

                # 修改调用方式，使用新的 RunnableSequence 方式
                full_response = chat_chain.invoke({"question": user_input})
                print(f"助手: {full_response}")
            except KeyboardInterrupt:
                # 处理Ctrl+C中断
                print("\n助手: 会话已中断")
                break
            except Exception as e:
                logger.error(f"聊天过程中出错: {str(e)}")
                print("\n助手: 抱歉，处理你的问题时出现了错误，请重试。")

    except Exception as e:
        logger.error(f"程序初始化失败: {str(e)}")


# 程序入口点
if __name__ == "__main__":
    # 设置环境变量，确保使用正确的线程数
    os.environ["OMP_NUM_THREADS"] = str(N_THREADS)
    # 禁用tokenizers并行化以避免警告
    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    # 启动主聊天程序
    cuda_chat_main()